import warnings
import numpy as np
from os import path, listdir
from scipy.optimize import fmin
from scipy.interpolate import splrep, splev
from csv import reader
from matplotlib import pyplot as plt
from math_functions import sg_filter
from pdb import set_trace
"""
Module to import, analyze, and plot data obtained 
from impact measurements
"""

# Global constants
g = 9.8             # gravity (m/s^2)

""" 
Structure to hold the measurement dtype for a numpy 
array.
"""
_type_def  = {
    'params' : {
        'names':['h','R','Es','vs','m'], 
        'formats':['f4']*5, 
        'titles':['acceleration height (m)','sphere radius (m)',
        'sphere modulus (Pa)','sphere Poissons ratio','sphere mass (kg)']
        },
    'est' : {
        'names':['Ep','vp','COR'], 
        'formats':['f4']*3, 
        'titles':['surface modulus (Pa)','surface Poissons ration', 
        'Coefficient of Restitution']
        },
    'dlc' :  {
        'names':['cal','mc','kc','mu'], 
        'formats':['f4']*4, 
        'titles':['dlc calibration (N/mV) or (m/mV)', 'crystal mass (kg)','crystal stiffness (N/m)', 
        'crystal damping factor (N s/m)']
        },
    'meta' : {
        'names':['name','desc', 'path'], 
        'formats':['|S40','|S100','|S80'],
        'titles':['file name','experiment description', 'file saved path'] 
        },
    'data' : {
        'names':['t','v'], 
        'formats':['1000f8','1000f8'],
        'titles':['time (seconds)','voltage (mV)']} 
    }
meas_descr = dict(names=[], formats=[], titles=[])
for k, v in _type_def.iteritems():
    meas_descr['names'] += v['names']
    meas_descr['formats'] += v['formats']
    meas_descr['titles'] += v['titles']
meas_dtype = np.dtype(meas_descr)

def import_file(fname):
    """
    Load a previously saved measurement file
    fname : string path name to text measurement file
    return numpy structured array with dtype meas_dtype
    """
    gui_vals = {}
    data = []
    lineNo = {'Begin Parameter':0,'End Parameter':0,
              'Begin Data':0, 'End Data':0}

    # Determine positions of data in file
    with open(fname, mode='r') as f:
        for ln, txt in enumerate(f):
            if txt.find("****Begin Parameter Values****") != -1:
                lineNo['Begin Parameter'] = ln
            elif txt.find("****End Parameter Values****") != -1: 
                lineNo['End Parameter'] = ln
            elif txt.find("****Begin Data****") != -1: 
                lineNo['Begin Data'] = ln + 1
            elif txt.find("****End Data****") != -1: 
                lineNo['End Data'] = ln
            
        # Read in the file data using the determined positions
        f.seek(0)
        csv_f = reader(f)
        for ln, txt in enumerate(csv_f):
            if (ln > lineNo['Begin Parameter'] and 
                ln < lineNo['End Parameter']):
                k, v = txt[0].split("=:=")
                try:
                    k, v = k.strip(), float(v.strip())
                except ValueError:
                    k, v = k.strip(), str(v.strip())
                gui_vals.update({k : v})
            elif (ln > lineNo['Begin Data'] and 
                  ln < lineNo['End Data']):
                data.append(txt)
    data = np.array(data,dtype='f8').transpose()

    # Create an array of the values
    meas = np.zeros((1), dtype=meas_dtype)
    meas['name'] = fname.split('/')[-1]
    meas['desc'] = gui_vals['ExpDescription']
    meas['path'] = gui_vals['SavePath']
    meas['h'] = gui_vals['DropHeight'] * 1.0e-2
    meas['m'] = gui_vals['SphereMass'] * 1.0e-3
    meas['R'] = gui_vals['SphereRadius'] * 1.0e-3
    meas['Es'] = gui_vals['SphereModulus'] * 1.0e9
    meas['vs' ] = gui_vals['SpherePR']
    meas['Ep'] = gui_vals['SurfaceModulus'] * 1.0e9
    meas['vp'] = gui_vals['SurfacePR']
    meas['COR'] = gui_vals['COR']
    meas['cal'] =  gui_vals['Calibration'] * 1.0e-9
    meas['mc'] = gui_vals['CrystalMass'] * 1e-3
    meas['kc'] = gui_vals['CrystalStiffness'] * 1e9
    meas['mu'] = gui_vals['CrystalDamping'] * 1e3
    meas['t'][0] =  data[0] * 1.0e-6
    meas['v'][0] =  data[1]

    return meas[0]

def import_folder(folder):
    """
    Attempt to import all the files in a director and
    create a numpy array contaiing the measurements.
    folder : sting path name to directory
    return numpy structured array with length as number of
        imported files and type meas_dtype
    """
    measurements = []
    for f in listdir(folder):
        fl_nm = path.join(folder,f)
        try:
            m = import_file(fl_nm)
        except KeyboardInterrupt:
            raise
        except:
            warnings.warn("Could not load file %s"%fl_nm)
        else:
            measurements.append(m)
    meas_array = np.zeros((len(measurements)), dtype=meas_dtype)
    for i, m in enumerate(measurements):
        meas_array[i] = m
    return meas_array

def plot_data(measurements):
    """
    Plot the measurements
    measurements - numpy structured array of measurement dtype
    """
    offset = [calculate_offset(t,v) for t,v in measurements['dat']]
    lines = [plt.plot(t-dt,v) for (t,v),dt in zip(measurements['dat'],offset)]

def crop_data(t, signal):
    """
    Crop the transient signal data to the portion
    where the sphere is in contact with the surface.
    t - numpy array time (s)
    signal - numpy array voltage (mV)
    return [nx1] numpy array of cropped signal
    """
    t_max = signal.argmax()
    v_max = np.max(signal)
    t_start = t_max - 5
    while np.average(signal[t_start:t_start+3]) > np.average(signal[t_start-3:t_start]) \
            and np.average(signal[t_start:t_start+3]) > 0.01*signal[t_max] \
            and t_start > 0:
        t_start = t_start - 1
    t_end = t_max + (np.sum(signal>v_max/2))/2
    while np.average(signal[t_end-3:t_end]) > np.average(signal[t_end:t_end+3]) \
            and np.average(signal[t_end-3:t_end]) > 0 \
            and t_end < t.shape[0]:
        t_end = t_end + 1
    return [t[t_start:t_end], signal[t_start:t_end]]
    
def calc_cor(t, f, m, h):
    """
    Calculate the coefficient of restitution from the force curve.
    t : [nx1] numpy array of time (s)
    f : [nx1] numpy array of transient force (Pa)
    m : mass of sphere (kg)
    h : height of acceleration (m)
    """
    integral = np.trapz(f,t)
    COR = 1./(m*np.sqrt(2.*g*h)) * integral - 1.
    return COR

def align_peak(t, v, t_f, v_f):
    """
    Align peaks of two curves.
    t : time axis of defining curve
    v : voltage of defining curve
    t_f : time of curve to be aligned
    v_f : voltage of curve to be aligned
    return t_f : new time axis of aligned curve
    """
    t_f -= - t_f[v_f.argmax()] + t[v.argmax()]
    return t_f

def calculate_offset(t, v):
    """
    Calculate the time offset of the signal, v, such that the impact begins
    at approximately 0 seconds.
    return offset time (seconds)
    TODD set the proper value for flag
    """
    flag = np.mean(v[:10]) + 3 * np.std(v[:10])
    flag = 100
    offset = next(i for i, n in enumerate(v) if n>flag)
    return t[offset]

def fit_parameters(t, v, f_theory, p_var, p_const):
    """
    t : time (s)
    v : voltge (mV)
    f_theory : theoretical function to fit to data (function must have
        all keyword arguments that are provided pby p_var, and p_const)
    p_var : dictionary key word arguments to f_theory that are varied
    p_const : dictionary key word arguments to f_theory that are constant
    return : dictionary of fit values and R2
    """
    t = t.copy()
    v = v.copy()
    def construct_func(f_theory, t, p_var, p_const):
        """Function reconstructs a new theoretical function, f_new, by substituting
        in keyword parameters from p_const, and converting keyword parameters,
        p_var to positional function arguments."""
        def f_new(*args):
            p_new = p_var.copy()
            for i, k in enumerate(p_var.keys()):
                """ The following statement and statement marked below
                are the key for proper assignment of keywords from tuple.
                Works because keys() returns same alphabetic order."""
                p_new[k] = args[i]
                t = np.linspace(0,10**-3,10**5)
                params = dict([('t',t)] + p_const.items() + p_new.items())
            return f_theory(**params)
        return f_new
    f_theory_mod = construct_func(f_theory, t, p_var, p_const)
            
    def objective(p, f_theory, t, v):
        results = f_theory(*p)
        t_fit = results['time']
        v_fit = results['voltage']
        t_fit, v_fit = crop_data(t_fit, v_fit)
        t_fit -= t_fit[np.argmax(v_fit)]

        spline = splrep(t_fit, v_fit)
        v_fit = splev(t, spline)
        residuals = np.sum((v**3 - v_fit**3)**2)
        return residuals

    t, v = crop_data(t, v)
    t -= t[np.argmax(v)]
    func_args = (f_theory_mod, t, v)
    """ The following statement and statement marked above
    are the key for proper assignment of keywords from tuple.
    Works because keys() returns same alphabetic order."""
    p_guess = np.array([p_var[k] for k in p_var.keys()])
    p_opt = fmin(objective, p_guess, args=func_args, xtol=10, ftol=10, 
                 maxiter=200, maxfun=50, full_output=1, disp=0, retall=0)

    p_fit = dict(zip(p_var.keys(),p_opt[0]))

    fit = f_theory_mod(*p_opt[0])
    t_fit = fit['time']
    v_fit = fit['voltage']
    t_fit, v_fit = crop_data(t_fit, v_fit)
    t_fit -= t_fit[np.argmax(v_fit)]
    spline = splrep(t_fit, v_fit)
    v_fit = splev(t, spline)
    SSres = np.sum((v - v_fit)**2)
    SStot = np.sum((v - np.mean(v))**2)
    p_fit['R2'] = 1 - SSres / SStot

    return p_fit
